import datetime

import baostock as bs
import pandas as pd

root_path = "../data/company-day/"
a_company_base_info_path = "../data/a_company_info.csv"


def get_a_base_ext(path=a_company_base_info_path):
    bs.login()
    # 显示登陆返回信息
    # 获取证券基本资料
    rs = bs.query_stock_basic(code="")
    # 打印结果集
    data_list = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(rs.get_row_data())
    result = pd.DataFrame(data_list, columns=rs.fields)
    ext = {}
    for i, row in result.iterrows():
        rs_profit = bs.query_profit_data(code=row['code'], year=2021, quarter=3)
        for field in rs_profit.fields:
            if not ext.__contains__(field):
                ext[field] = []

        if (rs_profit.error_code == '0') & rs_profit.next():
            data = rs_profit.get_row_data()
            for j in range(len(rs_profit.fields)):
                ext[rs_profit.fields[j]].append(data[j])
        else:
            for j in range(len(rs_profit.fields)):
                ext[rs_profit.fields[j]].append(0)
    for k, v in ext.items():
        try:
            if k == 'code' or k == 'pubDate' or k == 'statDate':
                continue
            result[k] = v
        except Exception as e:
            print(e)
    result.to_csv(path, encoding="utf-8", index=False)
    # 登出系统
    bs.logout()


def get_a_base_info(path=a_company_base_info_path, code=""):
    """
    获取A股数据基本信息
    :param path: 保存路径
    :param code: 股票编码
    :return:
    """
    # 登陆系统
    lg = bs.login()
    # 显示登陆返回信息
    print('login respond error_code:' + lg.error_code)
    print('login respond  error_msg:' + lg.error_msg)

    # 获取证券基本资料
    rs = bs.query_stock_basic(code=code)
    # rs = bs.query_stock_basic(code_name="浦发银行")  # 支持模糊查询
    print('query_stock_basic respond error_code:' + rs.error_code)
    print('query_stock_basic respond  error_msg:' + rs.error_msg)

    # 打印结果集
    data_list = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(rs.get_row_data())
    result = pd.DataFrame(data_list, columns=rs.fields)
    # 结果集输出到csv文件
    result.to_csv(path, encoding="utf-8", index=False)
    print(result)

    # 登出系统
    bs.logout()


def get_a_company_k_day_info(code="", name="", start_date="2020-10-01", end_date="", add=False):
    bs.login()
    # 显示登陆返回信息
    # 获取沪深A股历史K线数据
    # 详细指标参数，参见“历史行情指标参数”章节；“分钟线”参数与“日线”参数不同。“分钟线”不包含指数。
    # 分钟线指标：date,time,code,open,high,low,close,volume,amount,adjustflag
    # 周月线指标：date,code,open,high,low,close,volume,amount,adjustflag,turn,pctChg
    rs = bs.query_history_k_data_plus(code,
                                      "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,peTTM,pbMRQ,psTTM,pcfNcfTTM,isST",
                                      start_date=start_date, end_date=end_date,
                                      frequency="d", adjustflag="3")
    data_list = []
    while (rs.error_code == '0') & rs.next():
        # 获取一条记录，将记录合并在一起
        data_list.append(rs.get_row_data())
    result = pd.DataFrame(data_list, columns=rs.fields)

    if add:
        # 结果集输出到csv文件
        result.to_csv(root_path + code + "-" + name + ".csv", mode='a', index=False, header=None)
    else:
        result.to_csv(root_path + code + "-" + name + ".csv", index=False)

    # 登出系统
    bs.logout()


def download_k_day_info(path=a_company_base_info_path, start_date="2020-10-10", end_date="", add=False, start=0):
    """
    下载数据
    :param start_date: 开始时间
    :param end_date: 结束时间
    :return:
    """
    a_all_company = pd.read_csv(path)
    # 证券类型，其中1：股票，2：指数,3：其它
    a_all_company = a_all_company[a_all_company['type'] == 1]
    # 上市状态，其中1：上市，0：退市
    a_all_company = a_all_company[a_all_company['status'] == 1]
    # 获取全部的日K数据
    real_start = "2020-10-10"
    num = 0
    for row in a_all_company.iterrows():
        if add:
            real_start = start_date
        if num < start:
            num += 1
            continue
        get_a_company_k_day_info(row[1]['code'], row[1]['code_name'].replace("*", "-"), start_date=real_start,
                                 end_date=end_date, add=add)

        num += 1
        print(num)


if __name__ == '__main__':
    # download_k_day_info(start_date="2022-01-01", add=False, start=0)
    # download_k_day_info(start_date="2022-01-28", add=True)
    timeStr = datetime.date.today().strftime("%Y-%m-%d")
    download_k_day_info(start_date=timeStr, add=True, start=0)
    # download_k_day_info(start_date="2022-03-29", add=True, start=695)
    # get_a_base_ext()
    # get_a_base_info()
    today = datetime.date.today()
    print("over")
